import json
import re
import os
import random
from pathlib import Path

def process_jsonl_file(input_file, output_dir="output", sampled_dir="sampled_output", sample_size=10):
    """处理JSONL文件并按难度分割+随机抽样

    Args:
        input_file: 输入文件路径
        output_dir: 原始分割输出目录
        sampled_dir: 抽样输出目录
        sample_size: 每个难度抽样数量
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(sampled_dir, exist_ok=True)

    # 存储不同难度对应的数据
    difficulty_data = {}

    # 首先读取并分组所有数据
    with open(input_file, 'r', encoding='utf-8') as f:
        for line in f:
            try:
                data = json.loads(line)
                difficulty = data.get('difficulty', 'unknown').strip()
                cleaned_difficulty = re.sub(r'[^\w]+', '_', difficulty) or 'unknown'

                if cleaned_difficulty not in difficulty_data:
                    difficulty_data[cleaned_difficulty] = []
                difficulty_data[cleaned_difficulty].append(data)

            except json.JSONDecodeError:
                print(f"跳过无效行: {line.strip()}")

    # 处理原始分割
    for difficulty, items in difficulty_data.items():
        # 原始完整数据输出
        original_file = Path(output_dir) / f"problems_{difficulty}.json"
        with open(original_file, 'w', encoding='utf-8') as f:
            json.dump(items, f, indent=2, ensure_ascii=False)

        # 随机抽样处理
        sampled_file = Path(sampled_dir) / f"sampled_{difficulty}.json"
        sampled_items = random.sample(items, min(sample_size, len(items)))
        with open(sampled_file, 'w', encoding='utf-8') as f:
            json.dump(sampled_items, f, indent=2, ensure_ascii=False)

    print(f"处理完成:\n- 完整分割: {os.path.abspath(output_dir)}"
          f"\n- 随机抽样: {os.path.abspath(sampled_dir)} (每组最多{sample_size}个样本)")

if __name__ == "__main__":
    # 配置参数
    config = {
        "input_file": "test.jsonl",  # 输入文件
        "output_dir": "problems_full",          # 完整分割输出目录
        "sampled_dir": "problems_sampled",      # 抽样输出目录
        "sample_size": 3                       # 抽样数量
    }

    process_jsonl_file(**config)
